#coding:utf-8
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from ipdb import set_trace

class DetHeight(nn.Module):
    def __init__(self):
        super(DetHeight, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64, affine=True),
            nn.ReLU()
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 96, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(96, affine=True),
            nn.ReLU()
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(96, 128, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(128, affine=True),
            nn.ReLU()
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(128, 160, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(160, affine=True),
            nn.ReLU()
        )
        self.conv5 = nn.Sequential(
            nn.Conv2d(160, 192, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(192, affine=True),
            nn.ReLU()
        )
        self.conv6 = nn.Sequential(
            nn.Conv2d(192, 256, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(256, affine=True),
            nn.ReLU()
        )
        self.conv7 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(256, affine=True),
            nn.ReLU()
        )
        #self.pool = nn.AvgPool2d(kernel_size=(3, 5))
        self.fc1 = nn.Sequential(
            nn.Linear(256*3*5, 512),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(512, 1),
            #nn.ReLU()
        )

    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        x3 = self.conv3(x2)
        x4 = self.conv4(x3)
        x5 = self.conv5(x4)
        #set_trace()
        # x5 12 x 20
        x6 = self.conv6(x5)
        # x6 6 x 10
        x7 = self.conv7(x6)
        # x7 3 x 5
        x7 = x7.view(-1, 256*3*5)
        fc1 = self.fc1(x7)
        fc2 = self.fc2(fc1)
        fc2 = torch.squeeze(fc2)
        return fc2
        #x4 = x4.view(-1, )

    
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                n = m.weight.size(1)
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()
